import { flip, offset, shift, useFloating } from "@floating-ui/react";
import { useLexicalComposerContext } from "@lexical/react/LexicalComposerContext";
import {
  LexicalTypeaheadMenuPlugin,
  type MenuRenderFn,
} from "@lexical/react/LexicalTypeaheadMenuPlugin";
import { PRESET_VARIABLES } from "@shared/prompt-editor/constant";
import { $insertNodes, createCommand, type TextNode } from "lexical";
import { Fragment, useCallback, useMemo, useState } from "react";
import ReactDOM from "react-dom";
import { useTranslation } from "react-i18next";
import { triggerMatchForDoubleBrace } from "../../utils";
import { $createVariableValueNode } from "../nodes/variable-value-node";
import { PromptMenuItem } from "../prompt-menu-item";
import { VarOption } from "../var-option";

export const INSERT_QUERY_BLOCK_COMMAND = createCommand(
  "INSERT_QUERY_BLOCK_COMMAND",
);
export const DELETE_QUERY_BLOCK_COMMAND = createCommand(
  "DELETE_QUERY_BLOCK_COMMAND",
);

interface VariablePluginProps {
  isSystemPrompt: boolean;
}

export function VariablePlugin({ isSystemPrompt }: VariablePluginProps) {
  const { t } = useTranslation("translation", {
    keyPrefix: "chat.tool-bar.parameters.prompt-editor.menu-option",
  });

  const [editor] = useLexicalComposerContext();
  const triggerFn = triggerMatchForDoubleBrace();
  const { refs, floatingStyles, isPositioned } = useFloating({
    placement: "bottom-start",
    middleware: [
      offset(0), // fix hide cursor
      shift({
        padding: 8,
      }),
      flip(),
    ],
  });

  const [queryString, setQueryString] = useState<string | null>(null);

  const createOption = useCallback(
    (key: string, label: string, description: string) =>
      new VarOption({
        key,
        label,
        description,
        render: ({ isSelected, onSelect, onSetHighlight }) => (
          <PromptMenuItem
            title={label}
            isSelected={isSelected}
            onClick={onSelect}
            onMouseEnter={onSetHighlight}
          />
        ),
        onSelect: () =>
          editor.dispatchCommand(INSERT_QUERY_BLOCK_COMMAND, undefined),
      }),
    [editor],
  );

  const options = useMemo(() => {
    const descriptionMap = new Map([
      ["input", t("description.input")],
      ["date", t("description.date")],
      ["time", t("description.time")],
      ["datetime", t("description.datetime")],
      ["now", t("description.now")],
      ["model_id", t("description.model-id")],
    ]);
    return PRESET_VARIABLES.map((key) =>
      createOption(key, key, descriptionMap.get(key) || ""),
    ).filter((option) => {
      if (isSystemPrompt) return option.key !== "input";
      return true;
    });
  }, [createOption, t, isSystemPrompt]);

  const onSelectOption = useCallback(
    (
      selected: VarOption,
      nodeToRemove: TextNode | null,
      closeMenu: () => void,
    ) => {
      editor.update(() => {
        if (nodeToRemove) nodeToRemove.remove();
        $insertNodes([$createVariableValueNode(selected.key)]);
        closeMenu();
      });
    },
    [editor],
  );

  const menuRenderFn = useCallback<MenuRenderFn<VarOption>>(
    (
      anchorRef,
      { selectOptionAndCleanUp, setHighlightedIndex, selectedIndex },
    ) => {
      if (!anchorRef.current) return null;

      setTimeout(() => {
        if (anchorRef.current) refs.setReference(anchorRef.current);
      }, 0);

      return (
        <Fragment>
          {ReactDOM.createPortal(
            <div className="h-0 w-0">
              <div
                className="min-w-40 rounded-xl border bg-overlay p-1 text-overlay-fg shadow-lg"
                style={{
                  ...floatingStyles,
                  visibility: isPositioned ? "visible" : "hidden",
                }}
                ref={refs.setFloating}
              >
                {options.map((option, index) => (
                  <Fragment key={option.key}>
                    {option.renderMenuOption({
                      queryString,
                      isSelected: selectedIndex === index,
                      onSelect: () => {
                        selectOptionAndCleanUp(option);
                      },
                      onSetHighlight: () => {
                        setHighlightedIndex(index);
                      },
                    })}
                  </Fragment>
                ))}
              </div>
            </div>,
            anchorRef.current,
          )}
        </Fragment>
      );
    },
    [options, floatingStyles, isPositioned, refs, queryString],
  );

  return (
    <LexicalTypeaheadMenuPlugin<VarOption>
      triggerFn={triggerFn}
      onQueryChange={setQueryString}
      options={options}
      onSelectOption={onSelectOption}
      menuRenderFn={menuRenderFn}
      anchorClassName="z-[999999] translate-y-[calc(-100%-3px)]"
    />
  );
}
